#!/usr/bin/env python3
"""
Navigation Manager - 导航管理器
负责协调整个导航系统的运行
"""

import rclpy
from rclpy.node import Node
from rclpy.executors import MultiThreadedExecutor
from rclpy.callback_groups import ReentrantCallbackGroup

from geometry_msgs.msg import PoseStamped, Twist
from nav_msgs.msg import OccupancyGrid, Path
from sensor_msgs.msg import PointCloud2
from std_msgs.msg import Bool, String
from visualization_msgs.msg import Marker, MarkerArray

import numpy as np
import threading
import time
import traceback
from collections import deque
from datetime import datetime
from enum import Enum

class NavigationState(Enum):
    """导航状态枚举"""
    IDLE = "idle"
    PLANNING = "planning"
    NAVIGATING = "navigating"
    AVOIDING_OBSTACLE = "avoiding_obstacle"
    REACHED_GOAL = "reached_goal"
    FAILED = "failed"

class NavigationManager(Node):
    """
    导航管理器 - 系统核心协调节点
    功能：
    1. 接收目标点
    2. 协调路径规划
    3. 管理障碍物避障
    4. 控制无人机运动
    """
    
    def __init__(self):
        super().__init__('navigation_manager')
        
        # 初始化状态
        self.current_state = NavigationState.IDLE
        self.current_goal = None
        self.current_pose = None
        self.current_velocity = Twist()
        
        # 导航参数
        self.declare_parameter('goal_tolerance', 0.5)  # 目标容忍距离(米)
        self.declare_parameter('max_velocity', 2.0)    # 最大速度(m/s)
        self.declare_parameter('safety_distance', 1.0) # 安全距离(米)
        self.declare_parameter('path_lookahead_distance', 1.5)  # 路径前瞻距离
        self.declare_parameter('max_angular_velocity', 1.0)     # 最大角速度
        self.declare_parameter('emergency_stop_distance', 0.3)  # 紧急停止距离
        
        # 性能监控
        self.performance_log = deque(maxlen=1000)  # 保留最近1000次控制循环的性能数据
        self.last_performance_report = time.time()
        self.control_loop_count = 0
        self.error_count = 0
        
        # 状态历史记录
        self.state_history = deque(maxlen=100)
        self.last_state_change = time.time()
        
        self.goal_tolerance = self.get_parameter('goal_tolerance').value
        self.max_velocity = self.get_parameter('max_velocity').value
        self.safety_distance = self.get_parameter('safety_distance').value
        self.path_lookahead_distance = self.get_parameter('path_lookahead_distance').value
        self.max_angular_velocity = self.get_parameter('max_angular_velocity').value
        self.emergency_stop_distance = self.get_parameter('emergency_stop_distance').value
        
        # 回调组
        self.callback_group = ReentrantCallbackGroup()
        
        # 发布器
        self.cmd_vel_pub = self.create_publisher(
            Twist, '/cmd_vel', 10)
        self.status_pub = self.create_publisher(
            String, '/navigation_status', 10)
        self.path_marker_pub = self.create_publisher(
            MarkerArray, '/path_markers', 10)
            
        # 订阅器
        self.goal_sub = self.create_subscription(
            PoseStamped, '/goal_pose', self.goal_callback, 10,
            callback_group=self.callback_group)
        self.pose_sub = self.create_subscription(
            PoseStamped, '/current_pose', self.pose_callback, 10,
            callback_group=self.callback_group)
        self.obstacle_sub = self.create_subscription(
            Bool, '/obstacle_detected', self.obstacle_callback, 10,
            callback_group=self.callback_group)
        self.planned_path_sub = self.create_subscription(
            Path, '/planned_path', self.path_callback, 10,
            callback_group=self.callback_group)
            
        # 状态变量
        self.planned_path = None
        self.obstacle_detected = False
        self.path_index = 0
        
        # 创建定时器
        self.control_timer = self.create_timer(
            0.1, self.control_loop, callback_group=self.callback_group)  # 10Hz控制频率
            
        self.get_logger().info("导航管理器已启动")
        
    def goal_callback(self, msg):
        """接收目标点回调"""
        self.get_logger().info(f"收到新目标点: x={msg.pose.position.x:.2f}, y={msg.pose.position.y:.2f}, z={msg.pose.position.z:.2f}")
        self.current_goal = msg
        self.current_state = NavigationState.PLANNING
        self.path_index = 0
        self.planned_path = None
        
        # 发布状态
        status_msg = String()
        status_msg.data = f"开始导航到目标点: ({msg.pose.position.x:.2f}, {msg.pose.position.y:.2f}, {msg.pose.position.z:.2f})"
        self.status_pub.publish(status_msg)
        
    def pose_callback(self, msg):
        """当前位置回调"""
        self.current_pose = msg
        
    def obstacle_callback(self, msg):
        """障碍物检测回调"""
        self.obstacle_detected = msg.data
        if self.obstacle_detected and self.current_state == NavigationState.NAVIGATING:
            self.current_state = NavigationState.AVOIDING_OBSTACLE
            self.get_logger().warn("检测到障碍物，切换到避障模式")
            
    def path_callback(self, msg):
        """规划路径回调"""
        if len(msg.poses) > 0:
            self.planned_path = msg
            self.path_index = 0
            if self.current_state == NavigationState.PLANNING:
                self.current_state = NavigationState.NAVIGATING
                self.get_logger().info(f"收到规划路径，包含 {len(msg.poses)} 个路径点")
                self.visualize_path()
        else:
            self.get_logger().warn("收到空的规划路径")
            
    def control_loop(self):
        """主控制循环 - 增强版，包含性能监控和错误处理"""
        start_time = time.time()
        
        try:
            self.control_loop_count += 1
            
            # 基本检查
            if self.current_pose is None:
                if self.control_loop_count % 50 == 0:  # 每5秒报告一次
                    self.get_logger().warn("等待位置信息...")
                return
                
            if self.current_goal is None:
                if self.current_state != NavigationState.IDLE:
                    self.set_state(NavigationState.IDLE)
                return
            
            # 状态机处理
            if self.current_state == NavigationState.IDLE:
                self.stop_drone()
                
            elif self.current_state == NavigationState.PLANNING:
                self.stop_drone()
                # 规划超时检查
                if time.time() - self.last_state_change > 10.0:  # 10秒超时
                    self.get_logger().error("路径规划超时，切换到失败状态")
                    self.set_state(NavigationState.FAILED)
                
            elif self.current_state == NavigationState.NAVIGATING:
                if self.planned_path is None:
                    self.get_logger().warn("导航状态但无规划路径，重新规划")
                    self.set_state(NavigationState.PLANNING)
                    return
                    
                # 紧急停止检查
                if self.check_emergency_stop():
                    self.emergency_stop()
                    return
                    
                if self.obstacle_detected:
                    self.set_state(NavigationState.AVOIDING_OBSTACLE)
                    return
                    
                # 检查是否到达目标
                if self.is_goal_reached():
                    self.set_state(NavigationState.REACHED_GOAL)
                    self.get_logger().info("已到达目标点！")
                    status_msg = String()
                    status_msg.data = "目标到达"
                    self.status_pub.publish(status_msg)
                    return
                    
                # 路径跟踪
                self.follow_path()
                
            elif self.current_state == NavigationState.AVOIDING_OBSTACLE:
                if not self.obstacle_detected:
                    self.set_state(NavigationState.NAVIGATING)
                    self.get_logger().info("障碍物消失，恢复导航")
                else:
                    self.avoid_obstacle()
                    
            elif self.current_state == NavigationState.REACHED_GOAL:
                self.stop_drone()
                
            elif self.current_state == NavigationState.FAILED:
                self.stop_drone()
                
        except Exception as e:
            self.error_count += 1
            self.get_logger().error(f"控制循环错误 ({self.error_count}): {str(e)}")
            self.get_logger().error(f"错误详情: {traceback.format_exc()}")
            
            # 错误率过高时进入失败状态
            if self.error_count > 10:
                self.set_state(NavigationState.FAILED)
                
        finally:
            # 性能监控
            loop_time = time.time() - start_time
            self.performance_log.append({
                'timestamp': datetime.now(),
                'loop_time': loop_time,
                'state': self.current_state.value,
                'error_count': self.error_count
            })
            
            # 定期性能报告
            if time.time() - self.last_performance_report > 30.0:  # 每30秒
                self.report_performance()
                self.last_performance_report = time.time()
                
    def set_state(self, new_state):
        """状态切换管理"""
        if new_state != self.current_state:
            old_state = self.current_state
            self.current_state = new_state
            self.last_state_change = time.time()
            
            # 记录状态历史
            self.state_history.append({
                'timestamp': datetime.now(),
                'from_state': old_state.value,
                'to_state': new_state.value
            })
            
            self.get_logger().info(f"状态切换: {old_state.value} -> {new_state.value}")
            
    def check_emergency_stop(self):
        """紧急停止检查 - 优化版本"""
        if self.current_pose is None:
            return False
            
        # 这里可以添加更复杂的安全检查
        # 目前简化为障碍物检测
        return self.obstacle_detected
        
    def emergency_stop(self):
        """紧急停止"""
        self.get_logger().error("执行紧急停止！")
        
        # 立即停止
        stop_cmd = Twist()
        self.cmd_vel_pub.publish(stop_cmd)
        
        # 切换到失败状态
        self.set_state(NavigationState.FAILED)
        
        # 发布状态
        status_msg = String()
        status_msg.data = "紧急停止"
        self.status_pub.publish(status_msg)
        
    def report_performance(self):
        """性能报告 - 优化版本"""
        current_time = time.time()
        if current_time - self.last_performance_report > 10.0:  # 每10秒报告一次
            if self.control_loop_count > 0:
                avg_performance = sum(self.performance_log) / len(self.performance_log)
                self.get_logger().info(
                    f"性能统计 - 平均控制循环时间: {avg_performance*1000:.2f}ms, "
                    f"控制频率: {self.control_loop_count/10:.1f}Hz, "
                    f"错误率: {self.error_count/self.control_loop_count*100:.1f}%"
                )
            self.last_performance_report = current_time
            self.control_loop_count = 0
            self.error_count = 0
            
    def follow_path(self):
        """跟踪路径 - 优化版本"""
        if self.planned_path is None or len(self.planned_path.poses) == 0:
            self.stop_drone()
            return
            
        # 查找最佳跟踪点（前瞻算法优化）
        target_point = self.get_current_target_point_optimized()
        if target_point is None:
            self.get_logger().warn("无法获取目标跟踪点")
            self.stop_drone()
            return
            
        # 计算到目标点的向量
        current_pos = self.current_pose.pose.position
        dx = target_point.x - current_pos.x
        dy = target_point.y - current_pos.y
        dz = target_point.z - current_pos.z
        
        # 距离和角度计算
        distance_2d = (dx*dx + dy*dy) ** 0.5
        distance_3d = (dx*dx + dy*dy + dz*dz) ** 0.5
        
        if distance_3d < 0.1:  # 到达当前目标点
            self.path_index += 1
            return
            
        # 速度计算（动态调节）
        speed_factor = min(1.0, distance_3d / 2.0)  # 距离越近速度越慢
        
        cmd_vel = Twist()
        
        if distance_2d > 0.01:  # 避免除零
            # 线速度（归一化方向向量）
            cmd_vel.linear.x = (dx / distance_2d) * self.max_velocity * speed_factor
            cmd_vel.linear.y = (dy / distance_2d) * self.max_velocity * speed_factor
            
            # Z轴速度（独立控制）
            cmd_vel.linear.z = max(-self.max_velocity*0.5, 
                                  min(self.max_velocity*0.5, dz * speed_factor))
            
            # 角速度（面向目标方向）
            current_yaw = self.get_current_yaw()
            target_yaw = np.arctan2(dy, dx)
            yaw_error = target_yaw - current_yaw
            
            # 角度归一化到 [-π, π]
            while yaw_error > np.pi:
                yaw_error -= 2*np.pi
            while yaw_error < -np.pi:
                yaw_error += 2*np.pi
                
            cmd_vel.angular.z = max(-self.max_angular_velocity,
                                   min(self.max_angular_velocity, yaw_error * 2.0))
        
        # 安全限制
        self.apply_safety_limits(cmd_vel)
        self.cmd_vel_pub.publish(cmd_vel)
        
    def get_current_target_point_optimized(self):
        """获取当前目标跟踪点 - 优化版本"""
        if (self.planned_path is None or self.path_index >= len(self.planned_path.poses)):
            return None
            
        # 前瞻距离控制
        current_pos = self.current_pose.pose.position
        
        # 从当前索引开始搜索前瞻点
        best_index = self.path_index
        for i in range(self.path_index, len(self.planned_path.poses)):
            path_point = self.planned_path.poses[i].pose.position
            distance = ((path_point.x - current_pos.x)**2 + 
                       (path_point.y - current_pos.y)**2 + 
                       (path_point.z - current_pos.z)**2) ** 0.5
            
            if distance >= self.path_lookahead_distance:
                best_index = i
                break
            else:
                best_index = i  # 如果没有找到足够远的点，使用最后一个
                
        return self.planned_path.poses[best_index].pose.position
        
    def apply_safety_limits(self, cmd_vel):
        """应用安全限制 - 新增函数"""
        # 速度限制
        cmd_vel.linear.x = max(-self.max_velocity, min(self.max_velocity, cmd_vel.linear.x))
        cmd_vel.linear.y = max(-self.max_velocity, min(self.max_velocity, cmd_vel.linear.y))
        cmd_vel.linear.z = max(-self.max_velocity*0.5, min(self.max_velocity*0.5, cmd_vel.linear.z))
        cmd_vel.angular.z = max(-self.max_angular_velocity, min(self.max_angular_velocity, cmd_vel.angular.z))
        
    def avoid_obstacle(self):
        """避障控制 - 优化版本"""
        if not self.obstacle_detected:
            # 障碍物消失，恢复正常导航
            self.current_state = NavigationState.NAVIGATING
            return
            
        # 简单的避障策略：停止并稍微后退
        cmd_vel = Twist()
        cmd_vel.linear.x = -0.2  # 缓慢后退
        cmd_vel.linear.z = 0.3   # 稍微上升
        
        self.cmd_vel_pub.publish(cmd_vel)
        
        # 避障状态持续时间控制
        if time.time() - self.last_state_change > 3.0:  # 3秒后重新规划
            self.current_state = NavigationState.PLANNING
            
    def stop_drone(self):
        """停止无人机"""
        vel_cmd = Twist()  # 所有速度为0
        self.cmd_vel_pub.publish(vel_cmd)
        
    def is_goal_reached(self):
        """检查是否到达目标"""
        if self.current_goal is None or self.current_pose is None:
            return False
            
        dx = self.current_goal.pose.position.x - self.current_pose.pose.position.x
        dy = self.current_goal.pose.position.y - self.current_pose.pose.position.y
        dz = self.current_goal.pose.position.z - self.current_pose.pose.position.z
        
        distance = np.sqrt(dx**2 + dy**2 + dz**2)
        return distance < self.goal_tolerance
        
    def get_current_yaw(self):
        """从四元数获取当前偏航角"""
        if self.current_pose is None:
            return 0.0
            
        # 从四元数计算偏航角
        q = self.current_pose.pose.orientation
        yaw = np.arctan2(
            2.0 * (q.w * q.z + q.x * q.y),
            1.0 - 2.0 * (q.y**2 + q.z**2)
        )
        return yaw
        
    def visualize_path(self):
        """可视化规划路径"""
        if self.planned_path is None:
            return
            
        marker_array = MarkerArray()
        
        # 路径线条
        line_marker = Marker()
        line_marker.header.frame_id = "map"
        line_marker.header.stamp = self.get_clock().now().to_msg()
        line_marker.ns = "planned_path"
        line_marker.id = 0
        line_marker.type = Marker.LINE_STRIP
        line_marker.action = Marker.ADD
        line_marker.scale.x = 0.1
        line_marker.color.r = 0.0
        line_marker.color.g = 1.0
        line_marker.color.b = 0.0
        line_marker.color.a = 1.0
        
        for pose in self.planned_path.poses:
            line_marker.points.append(pose.pose.position)
            
        marker_array.markers.append(line_marker)
        
        # 目标点标记
        if self.current_goal is not None:
            goal_marker = Marker()
            goal_marker.header.frame_id = "map"
            goal_marker.header.stamp = self.get_clock().now().to_msg()
            goal_marker.ns = "goal"
            goal_marker.id = 1
            goal_marker.type = Marker.SPHERE
            goal_marker.action = Marker.ADD
            goal_marker.pose = self.current_goal.pose
            goal_marker.scale.x = 0.5
            goal_marker.scale.y = 0.5
            goal_marker.scale.z = 0.5
            goal_marker.color.r = 1.0
            goal_marker.color.g = 0.0
            goal_marker.color.b = 0.0
            goal_marker.color.a = 1.0
            
            marker_array.markers.append(goal_marker)
            
        self.path_marker_pub.publish(marker_array)

def main(args=None):
    rclpy.init(args=args)
    
    navigation_manager = NavigationManager()
    
    # 使用多线程执行器
    executor = MultiThreadedExecutor()
    
    try:
        rclpy.spin(navigation_manager, executor=executor)
    except KeyboardInterrupt:
        pass
    finally:
        navigation_manager.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main() 